import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist

# 模拟参数
SNR_dB = np.arange(0, 21, 2)  # 信噪比范围（dB）
num_bits = 1000  # 数据长度（位）
mod_order_bpsk = 2  # BPSK调制
mod_order_qpsk = 4  # QPSK调制
mod_order_16qam = 16  # 16-QAM调制


# BPSK调制
def bpsk_modulate(data):
    return 2 * data - 1  # BPSK映射：0 -> -1，1 -> 1


# QPSK调制
def qpsk_modulate(data):
    return np.array([np.cos(np.pi / 2 * i) + 1j * np.sin(np.pi / 2 * i) for i in data])


# 16-QAM调制
def qam16_modulate(data):
    mapping = {
        0: (-3 - 3j), 1: (-3 - 1j), 2: (-3 + 3j), 3: (-3 + 1j),
        4: (-1 - 3j), 5: (-1 - 1j), 6: (-1 + 3j), 7: (-1 + 1j),
        8: (3 - 3j), 9: (3 - 1j), 10: (3 + 3j), 11: (3 + 1j),
        12: (1 - 3j), 13: (1 - 1j), 14: (1 + 3j), 15: (1 + 1j)
    }
    return np.array([mapping[i] for i in data])


# AWGN信道
def awgn_channel(signal, snr_db):
    """模拟AWGN信道"""
    snr = 10 ** (snr_db / 10)  # 线性信噪比
    signal_power = np.mean(np.abs(signal) ** 2)
    noise_power = signal_power / snr
    noise = np.sqrt(noise_power) * np.random.randn(len(signal))  # 生成AWGN噪声
    received_signal = signal + noise
    return received_signal


# 解调BPSK信号
def bpsk_demodulation(received_signal):
    """BPSK解调"""
    return np.where(received_signal > 0, 1, 0)


# 解调QPSK信号
def qpsk_demodulation(received_signal):
    """QPSK解调"""
    phase = np.angle(received_signal)
    return np.array(np.round((phase + np.pi) / (np.pi / 2)), dtype=int)


# 16-QAM解调
def qam16_demodulation(received_signal):
    """16-QAM解调：将接收到的信号映射到最接近的QAM符号"""
    mapping = {
        (-3 - 3j): 0, (-3 - 1j): 1, (-3 + 3j): 2, (-3 + 1j): 3,
        (-1 - 3j): 4, (-1 - 1j): 5, (-1 + 3j): 6, (-1 + 1j): 7,
        (3 - 3j): 8, (3 - 1j): 9, (3 + 3j): 10, (3 + 1j): 11,
        (1 - 3j): 12, (1 - 1j): 13, (1 + 3j): 14, (1 + 1j): 15
    }

    symbols = np.array(list(mapping.keys()))  # 定义QAM16符号
    # 使用np.abs计算每个符号的距离（即模长），避免使用复数类型的距离计算
    distances = cdist(np.abs(received_signal.reshape(-1, 1)), np.abs(symbols.reshape(-1, 1)), metric='euclidean')
    closest_indices = np.argmin(distances, axis=1)  # 找到距离最小的符号
    return np.array([mapping[symbols[i]] for i in closest_indices])


# 计算误码率（BER）
def calculate_ber(original_data, demodulated_data):
    """计算误码率（BER）"""
    return np.sum(original_data != demodulated_data) / len(original_data)


# 仿真不同SNR下的BER
ber_bpsk = []
ber_qpsk = []
ber_qam16 = []
for snr in SNR_dB:
    # 生成数据
    data_bpsk = np.random.randint(0, mod_order_bpsk, num_bits)
    data_qpsk = np.random.randint(0, mod_order_qpsk, num_bits // 2)
    data_16qam = np.random.randint(0, mod_order_16qam, num_bits // 4)

    # 调制
    bpsk_signal = bpsk_modulate(data_bpsk)
    qpsk_signal = qpsk_modulate(data_qpsk)
    qam16_signal = qam16_modulate(data_16qam)

    # 通过AWGN信道传输
    received_signal_bpsk = awgn_channel(bpsk_signal, snr)
    received_signal_qpsk = awgn_channel(qpsk_signal, snr)
    received_signal_qam16 = awgn_channel(qam16_signal, snr)

    # 解调
    demodulated_bpsk = bpsk_demodulation(received_signal_bpsk)
    demodulated_qpsk = qpsk_demodulation(received_signal_qpsk)
    demodulated_qam16 = qam16_demodulation(received_signal_qam16)

    # 计算BER
    ber_bpsk.append(calculate_ber(data_bpsk, demodulated_bpsk))
    ber_qpsk.append(calculate_ber(data_qpsk, demodulated_qpsk))
    ber_qam16.append(calculate_ber(data_16qam, demodulated_qam16))

# 绘制多个图表
# 1. 绘制BPSK, QPSK 和 16-QAM的BER vs SNR图
plt.figure(figsize=(12, 10))

# BPSK BER 图
plt.figure(1)
plt.semilogy(SNR_dB, ber_bpsk, label="BPSK")
plt.title("BPSK: Bit Error Rate (BER) vs Signal-to-Noise Ratio (SNR)")
plt.xlabel("SNR (dB)")
plt.ylabel("BER")
plt.grid(True)
plt.legend()

# QPSK BER 图
plt.figure(2)
plt.semilogy(SNR_dB, ber_qpsk, label="QPSK")
plt.title("QPSK: Bit Error Rate (BER) vs Signal-to-Noise Ratio (SNR)")
plt.xlabel("SNR (dB)")
plt.ylabel("BER")
plt.grid(True)
plt.legend()

# 16-QAM BER 图
plt.figure(3)
plt.semilogy(SNR_dB, ber_qam16, label="16-QAM")
plt.title("16-QAM: Bit Error Rate (BER) vs Signal-to-Noise Ratio (SNR)")
plt.xlabel("SNR (dB)")
plt.ylabel("BER")
plt.grid(True)
plt.legend()

# 2. 绘制星座图
# BPSK星座图
plt.figure(4)
plt.scatter(np.real(bpsk_signal), np.imag(bpsk_signal), color='blue', label='BPSK')
plt.title("BPSK Constellation")
plt.xlabel("Real")
plt.ylabel("Imaginary")
plt.grid(True)
plt.legend()

# QPSK星座图
plt.figure(5)
plt.scatter(np.real(qpsk_signal), np.imag(qpsk_signal), color='red', label='QPSK')
plt.title("QPSK Constellation")
plt.xlabel("Real")
plt.ylabel("Imaginary")
plt.grid(True)
plt.legend()

# 16-QAM星座图
plt.figure(6)
plt.scatter(np.real(qam16_signal), np.imag(qam16_signal), color='green', label='16-QAM')
plt.title("16-QAM Constellation")
plt.xlabel("Real")
plt.ylabel("Imaginary")
plt.grid(True)
plt.legend()

# 3. 信号与噪声的对比
plt.figure(7)
plt.plot(np.real(bpsk_signal), label="Original BPSK Signal")
plt.plot(np.real(received_signal_bpsk), label="Received BPSK Signal with Noise", linestyle="--")
plt.title("BPSK Signal and Noisy Signal Comparison")
plt.xlabel("Sample Index")
plt.ylabel("Amplitude")
plt.legend()
plt.grid(True)

# 显示所有图形
plt.show()

# 输出不同调制方式的BER结果
for snr, bpsk, qpsk, qam16 in zip(SNR_dB, ber_bpsk, ber_qpsk, ber_qam16):
    print(f"SNR = {snr} dB, BPSK BER = {bpsk}, QPSK BER = {qpsk}, 16-QAM BER = {qam16}")
